# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/4 下午5:45
@Author yzpang
@Function: 数据集与数据加载器构建
"""

from .util import parse_labels, read_local_dataset, preprocess_function
from modelserver.configs.server_config import MULTI_CLASS
import functools
from paddlenlp.data import DataCollatorWithPadding
from paddle.io import BatchSampler, DataLoader
from paddlenlp.datasets import load_dataset


def dataset(tokenizer, label_path, example_path, classification_type=MULTI_CLASS):

    """构建数据集"""
    id2label, label2id = parse_labels(label_path)
    example_ds = load_dataset(read_local_dataset,
                              path=example_path,
                              label2id=label2id,
                              classification_type=classification_type,
                              lazy=False)
    example_func = functools.partial(preprocess_function,
                                     tokenizer=tokenizer,
                                     max_length=128,
                                     classification_type=classification_type,
                                     labels_nums=len(label2id))
    example_ds = example_ds.map(example_func)
    return example_ds


def dataloader(tokenizer, label_path, example_path, classification_type=MULTI_CLASS):

    """构建数据加载器"""
    example_ds = dataset(tokenizer, label_path, example_path, classification_type)
    collate_fn = DataCollatorWithPadding(tokenizer)
    batch_sampler = BatchSampler(example_ds, batch_size=32, shuffle=True)
    data_loader = DataLoader(dataset=example_ds, batch_sampler=batch_sampler, collate_fn=collate_fn)
    return data_loader


def label_list(label_path):
    """返回类别文件中类别数量"""
    return parse_labels(label_path)
